{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Movie Recommendation System using Graph-based Methods\n",
    "This notebook implements a complete movie recommendation system using graph-based approaches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required libraries\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from datetime import datetime, timedelta\n",
    "from collections import defaultdict\n",
    "import random\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Global variables\n",
    "GENRE_LIST = ['Action', 'Drama', 'Comedy', 'Romance', 'Thriller']\n",
    "movie_data = {}\n",
    "ACCUMULATION_STEPS = 4  # Define accumulation steps constant"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph Building and Feature Engineering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    "from datetime import datetime\n",
    "\n",
    "def build_movie_graph(interactions):\n",
    "    G = nx.Graph()\n",
    "    for user_id, movie_id, rating, timestamp in interactions:\n",
    "        # Add user node with type attribute\n",
    "        user_node = f\"user_{user_id}\"\n",
    "        G.add_node(user_node, type=\"user\", last_active=timestamp)\n",
    "        \n",
    "        # Add movie node with type attribute\n",
    "        movie_node = f\"movie_{movie_id}\"\n",
    "        G.add_node(movie_node, type=\"movie\")\n",
    "        \n",
    "        # Add or update edge\n",
    "        if G.has_edge(user_node, movie_node):\n",
    "            G[user_node][movie_node]['weight'] += 1\n",
    "            G[user_node][movie_node]['last_interaction'] = timestamp\n",
    "            G[user_node][movie_node]['ratings'].append(float(rating))\n",
    "        else:\n",
    "            G.add_edge(user_node, movie_node, \n",
    "                      weight=1, \n",
    "                      last_interaction=timestamp, \n",
    "                      ratings=[float(rating)])\n",
    "    return G\n",
    "\n",
    "def normalize_edge_weights(G):\n",
    "    for edge in G.edges():\n",
    "        ratings = G.edges[edge].get('ratings', [])\n",
    "        if ratings:\n",
    "            avg_rating = np.mean(ratings)\n",
    "            G.edges[edge]['rating'] = avg_rating  # Store average rating\n",
    "            G.edges[edge]['weight'] = avg_rating / 5.0  # Normalize to [0,1]\n",
    "        else:\n",
    "            G.edges[edge]['rating'] = 0\n",
    "            G.edges[edge]['weight'] = 0\n",
    "\n",
    "\n",
    "\n",
    "def engineer_user_features(G, user_data):\n",
    "    user_features = {}\n",
    "    for node, data in G.nodes(data=True):\n",
    "        if data['type'] == 'user':\n",
    "            user_id = node.split('_')[1]\n",
    "            if user_id in user_data:  # Add check if user exists in user_data\n",
    "                user_info = user_data[user_id]\n",
    "                # Basic features\n",
    "                features = [\n",
    "                    float(user_info['age']),  # Convert to float\n",
    "                    float(user_info['gender_encoded']),\n",
    "                    float(user_info['location_encoded']),\n",
    "                ]\n",
    "                \n",
    "                # Behavioral features\n",
    "                # Get ratings directly from edges\n",
    "                ratings = [G[node][edge].get('rating', 0) for edge in G[node]]\n",
    "                avg_rating = np.mean(ratings) if ratings else 0\n",
    "                rating_count = G.degree(node)\n",
    "                features.extend([avg_rating, float(rating_count)])\n",
    "                \n",
    "                # Add genre preferences\n",
    "                genre_preferences = calculate_genre_preferences(G, node)\n",
    "                features.extend(genre_preferences)\n",
    "                \n",
    "                user_features[node] = np.array(features, dtype=np.float32)\n",
    "    \n",
    "    return user_features\n",
    "\n",
    "def calculate_genre_preferences(G, user_node):\n",
    "    genre_counts = {genre: 0 for genre in GENRE_LIST}\n",
    "    total_ratings = 0\n",
    "    \n",
    "    # Iterate through neighboring movie nodes\n",
    "    for neighbor in G[user_node]:\n",
    "        if neighbor.startswith('movie_'):\n",
    "            movie_id = neighbor.split('_')[1]\n",
    "            if movie_id in movie_data:\n",
    "                genres = movie_data[movie_id]['genres']\n",
    "                rating = G[user_node][neighbor]['rating']\n",
    "                for genre in genres:\n",
    "                    if genre in genre_counts:\n",
    "                        genre_counts[genre] += rating\n",
    "                total_ratings += 1\n",
    "    \n",
    "    # Normalize genre preferences\n",
    "    genre_preferences = []\n",
    "    for genre in GENRE_LIST:\n",
    "        if total_ratings > 0:\n",
    "            genre_preferences.append(genre_counts[genre] / total_ratings)\n",
    "        else:\n",
    "            genre_preferences.append(0)\n",
    "    \n",
    "    return genre_preferences\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loss Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bce_loss(predictions, targets):\n",
    "    return F.binary_cross_entropy_with_logits(predictions, targets)\n",
    "\n",
    "def mse_loss(predictions, targets):\n",
    "    return F.mse_loss(predictions, targets)\n",
    "\n",
    "def bpr_loss(pos_scores, neg_scores):\n",
    "    return -F.logsigmoid(pos_scores - neg_scores).mean()\n",
    "\n",
    "def margin_ranking_loss(pos_scores, neg_scores, margin=0.5):\n",
    "    return F.margin_ranking_loss(pos_scores, neg_scores, torch.ones_like(pos_scores), margin=margin)\n",
    "\n",
    "def combined_loss(pred_ratings, true_ratings, pos_scores, neg_scores, alpha=0.5):\n",
    "    # Convert inputs to floating point tensors if they aren't already\n",
    "    pred_ratings = pred_ratings.float()\n",
    "    true_ratings = true_ratings.float()\n",
    "    pos_scores = pos_scores.float()\n",
    "    neg_scores = neg_scores.float()\n",
    "    \n",
    "    rating_loss = mse_loss(pred_ratings, true_ratings)\n",
    "    ranking_loss = bpr_loss(pos_scores, neg_scores)\n",
    "    return alpha * rating_loss + (1 - alpha) * ranking_loss\n",
    "\n",
    "\n",
    "def evaluate_model(model, val_graph):\n",
    "    \"\"\"\n",
    "    Evaluate the model on validation graph and return validation loss\n",
    "    \n",
    "    Args:\n",
    "        model: The GNN model\n",
    "        val_graph: Validation graph (NetworkX graph)\n",
    "        \n",
    "    Returns:\n",
    "        float: Average validation loss\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    total_loss = 0\n",
    "    batch_count = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for batch in generate_batches(val_graph, batch_size=32):\n",
    "            user_nodes, pos_movie_nodes, neg_movie_nodes, ratings = batch\n",
    "            \n",
    "            # Get model predictions\n",
    "            pred_ratings, pos_scores, neg_scores = model(\n",
    "                val_graph, \n",
    "                user_nodes, \n",
    "                pos_movie_nodes, \n",
    "                neg_movie_nodes\n",
    "            )\n",
    "            \n",
    "            # Calculate combined loss\n",
    "            loss = combined_loss(\n",
    "                pred_ratings,\n",
    "                ratings,\n",
    "                pos_scores,\n",
    "                neg_scores\n",
    "            )\n",
    "            \n",
    "            total_loss += loss.item()\n",
    "            batch_count += 1\n",
    "    \n",
    "    # Return average loss\n",
    "    return total_loss / max(1, batch_count)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_batches(graph, batch_size=2):\n",
    "    for _ in range(3):  # Generate 3 dummy batches\n",
    "        subgraph, user_nodes, pos_movies, neg_movies = create_mini_batch(\n",
    "            graph, batch_size=batch_size, n_pos=2, n_neg=2\n",
    "        )\n",
    "        # Create ratings tensor with gradient computation enabled\n",
    "        ratings = torch.rand(batch_size, 1, requires_grad=True)\n",
    "        yield user_nodes, pos_movies, neg_movies, ratings\n",
    "\n",
    "def train_model(model, train_graph, val_graph, optimizer, scheduler, epochs, patience=5):\n",
    "    best_val_loss = float('inf')\n",
    "    no_improve_count = 0\n",
    "    \n",
    "    for epoch in range(epochs):\n",
    "        model.train()\n",
    "        total_loss = 0\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        for batch_idx, batch in enumerate(generate_batches(train_graph)):\n",
    "            user_nodes, pos_movie_nodes, neg_movie_nodes, ratings = batch\n",
    "            pred_ratings, pos_scores, neg_scores = model(train_graph, user_nodes, pos_movie_nodes, neg_movie_nodes)\n",
    "            \n",
    "            loss = combined_loss(pred_ratings, ratings, pos_scores, neg_scores)\n",
    "            loss = loss / ACCUMULATION_STEPS\n",
    "            loss.backward()\n",
    "            \n",
    "            if (batch_idx + 1) % ACCUMULATION_STEPS == 0:\n",
    "                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n",
    "                optimizer.step()\n",
    "                optimizer.zero_grad()\n",
    "                \n",
    "            total_loss += loss.item() * ACCUMULATION_STEPS\n",
    "            \n",
    "        avg_train_loss = total_loss / len(train_graph)\n",
    "        val_loss = evaluate_model(model, val_graph)\n",
    "        \n",
    "        print(f\"Epoch {epoch}, Train Loss: {avg_train_loss:.4f}, Val Loss: {val_loss:.4f}\")\n",
    "        scheduler.step(val_loss)\n",
    "        \n",
    "        if val_loss < best_val_loss:\n",
    "            best_val_loss = val_loss\n",
    "            no_improve_count = 0\n",
    "            torch.save(model.state_dict(), 'best_model.pth')\n",
    "        else:\n",
    "            no_improve_count += 1\n",
    "            if no_improve_count >= patience:\n",
    "                print(f\"Early stopping after {epoch} epochs\")\n",
    "                break\n",
    "                \n",
    "    model.load_state_dict(torch.load('best_model.pth'))\n",
    "    return model\n",
    "\n",
    "\n",
    "def sample_neighbors(graph, nodes, n_neighbors, n_hops):\n",
    "    sampled_nodes = set(nodes)\n",
    "    for _ in range(n_hops):\n",
    "        new_nodes = set()\n",
    "        for node in sampled_nodes:\n",
    "            neighbors = list(graph.neighbors(node))\n",
    "            sampled = random.sample(neighbors, min(n_neighbors, len(neighbors)))\n",
    "            new_nodes.update(sampled)\n",
    "        sampled_nodes.update(new_nodes)\n",
    "    return list(sampled_nodes)\n",
    "\n",
    "def create_mini_batch(G, batch_size, n_pos=5, n_neg=5, n_neighbors=10, n_hops=2):\n",
    "    # Get all user nodes\n",
    "    all_user_nodes = [n for n in G.nodes() if n.startswith('user_')]\n",
    "    if not all_user_nodes:\n",
    "        raise ValueError(\"No user nodes found in graph\")\n",
    "    \n",
    "    # Randomly select users\n",
    "    user_nodes = random.sample(all_user_nodes, min(batch_size, len(all_user_nodes)))\n",
    "    \n",
    "    # Create subgraph\n",
    "    subgraph = G.subgraph(user_nodes).copy()\n",
    "    \n",
    "    # Sample positive and negative movies\n",
    "    pos_movies = []\n",
    "    neg_movies = []\n",
    "    all_movies = [n for n in G.nodes() if n.startswith('movie_')]\n",
    "    \n",
    "    for user in user_nodes:\n",
    "        user_movies = [n for n in G[user] if n.startswith('movie_')]\n",
    "        pos_sample = random.sample(user_movies, min(n_pos, len(user_movies))) if user_movies else []\n",
    "        available_neg = list(set(all_movies) - set(user_movies))\n",
    "        neg_sample = random.sample(available_neg, min(n_neg, len(available_neg))) if available_neg else []\n",
    "        \n",
    "        pos_movies.extend(pos_sample)\n",
    "        neg_movies.extend(neg_sample)\n",
    "    \n",
    "    return subgraph, user_nodes, pos_movies, neg_movies\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cold Start Solutions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "def initialize_new_movie(movie_id, similar_movies, movie_embeddings):\n",
    "    if not similar_movies:\n",
    "        return np.zeros(next(iter(movie_embeddings.values())).shape)\n",
    "    \n",
    "    similar_embeddings = [movie_embeddings[m] for m in similar_movies if m in movie_embeddings]\n",
    "    if not similar_embeddings:\n",
    "        return np.zeros(next(iter(movie_embeddings.values())).shape)\n",
    "    \n",
    "    return np.mean(similar_embeddings, axis=0)\n",
    "\n",
    "def create_initial_edges(user_id, preferences, movie_graph):\n",
    "    for pref in preferences:\n",
    "        similar_movies = [m for m in movie_graph.nodes if pref in movie_graph.nodes[m]['attributes']]\n",
    "        for movie in similar_movies:\n",
    "            movie_graph.add_edge(user_id, movie, weight=0.5)\n",
    "    return movie_graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main Execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting recommendation system test...\n",
      "\n",
      "1. Sample data created\n",
      "\n",
      "2. Graph built successfully\n",
      "Number of nodes: 6\n",
      "Number of edges: 5\n",
      "\n",
      "3. Edge weights normalized\n",
      "Sample edge weights: ('user_1', 'movie_101', {'weight': 0.9, 'last_interaction': 1732915204.994279, 'ratings': [4.5], 'rating': 4.5})\n",
      "\n",
      "4. User features engineered\n",
      "Number of users with features: 3\n",
      "Feature vector size: 10\n",
      "\n",
      "5. Mini-batch created\n",
      "Subgraph nodes: 2\n",
      "Sample users: ['user_2', 'user_3']\n",
      "\n",
      "6. Loss calculation tested\n",
      "Combined loss value: 1.4735026359558105\n",
      "\n",
      "7. Cold start solution tested\n",
      "New movie embedding: [0.15 0.25 0.35]\n"
     ]
    }
   ],
   "source": [
    "import networkx as nx\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import random\n",
    "from datetime import datetime\n",
    "\n",
    "# Sample Data Creation\n",
    "\n",
    "def create_sample_data():\n",
    "    # Sample interactions data\n",
    "    interactions = [\n",
    "        (1, 101, 4.5, datetime.now().timestamp()),\n",
    "        (1, 102, 3.5, datetime.now().timestamp()),\n",
    "        (2, 101, 5.0, datetime.now().timestamp()),\n",
    "        (2, 103, 4.0, datetime.now().timestamp()),\n",
    "        (3, 102, 3.0, datetime.now().timestamp()),\n",
    "    ]\n",
    "    \n",
    "    # Sample user data\n",
    "    user_data = {\n",
    "        '1': {'age': 25, 'gender_encoded': 1, 'location_encoded': 2},\n",
    "        '2': {'age': 30, 'gender_encoded': 0, 'location_encoded': 1},\n",
    "        '3': {'age': 35, 'gender_encoded': 1, 'location_encoded': 3},\n",
    "    }\n",
    "    \n",
    "    # Sample movie data\n",
    "    global movie_data\n",
    "    movie_data = {\n",
    "        '101': {'genres': ['Action', 'Drama']},\n",
    "        '102': {'genres': ['Comedy', 'Romance']},\n",
    "        '103': {'genres': ['Action', 'Thriller']},\n",
    "    }\n",
    "    \n",
    "    return interactions, user_data\n",
    "\n",
    "# Simple Graph Neural Network Model for Testing\n",
    "class SimpleGNN(nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dim):\n",
    "        super(SimpleGNN, self).__init__()\n",
    "        self.layer1 = nn.Linear(input_dim, hidden_dim)\n",
    "        self.layer2 = nn.Linear(hidden_dim, 1)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.layer1(x))\n",
    "        return self.layer2(x)\n",
    "\n",
    "# Test function\n",
    "def test_recommendation_system():\n",
    "    print(\"Starting recommendation system test...\")\n",
    "    \n",
    "    # 1. Create sample data\n",
    "    interactions, user_data = create_sample_data()\n",
    "    print(\"\\n1. Sample data created\")\n",
    "    \n",
    "    # 2. Test graph building\n",
    "    G = build_movie_graph(interactions)\n",
    "    print(\"\\n2. Graph built successfully\")\n",
    "    print(f\"Number of nodes: {G.number_of_nodes()}\")\n",
    "    print(f\"Number of edges: {G.number_of_edges()}\")\n",
    "    \n",
    "    # 3. Test edge weight normalization\n",
    "    normalize_edge_weights(G)\n",
    "    print(\"\\n3. Edge weights normalized\")\n",
    "    print(\"Sample edge weights:\", next(iter(G.edges(data=True))))\n",
    "    \n",
    "    # 4. Test feature engineering\n",
    "    user_features = engineer_user_features(G, user_data)\n",
    "    print(\"\\n4. User features engineered\")\n",
    "    print(f\"Number of users with features: {len(user_features)}\")\n",
    "    if user_features:\n",
    "        print(f\"Feature vector size: {len(next(iter(user_features.values())))}\")\n",
    "    \n",
    "    # 5. Test batch creation\n",
    "    try:\n",
    "        subgraph, user_nodes, pos_movies, neg_movies = create_mini_batch(\n",
    "            G, batch_size=2, n_pos=2, n_neg=2, n_neighbors=2, n_hops=1\n",
    "        )\n",
    "        print(\"\\n5. Mini-batch created\")\n",
    "        print(f\"Subgraph nodes: {subgraph.number_of_nodes()}\")\n",
    "        print(f\"Sample users: {user_nodes[:2]}\")\n",
    "    except Exception as e:\n",
    "        print(f\"\\n5. Mini-batch creation failed: {str(e)}\")\n",
    "    \n",
    "    # 6. Test model and loss functions\n",
    "    if user_features:\n",
    "        feature_size = len(next(iter(user_features.values())))\n",
    "        model = SimpleGNN(input_dim=feature_size, hidden_dim=32)\n",
    "        sample_predictions = torch.randn(5)\n",
    "        sample_targets = torch.randn(5)\n",
    "        sample_pos_scores = torch.randn(5)\n",
    "        sample_neg_scores = torch.randn(5)\n",
    "        \n",
    "        loss = combined_loss(\n",
    "            sample_predictions, \n",
    "            sample_targets,\n",
    "            sample_pos_scores, \n",
    "            sample_neg_scores\n",
    "        )\n",
    "        print(\"\\n6. Loss calculation tested\")\n",
    "        print(f\"Combined loss value: {loss.item()}\")\n",
    "    \n",
    "    # 7. Test cold start solutions\n",
    "    new_movie_id = '104'\n",
    "    similar_movies = ['101', '102']\n",
    "    dummy_embeddings = {\n",
    "        '101': np.array([0.1, 0.2, 0.3]),\n",
    "        '102': np.array([0.2, 0.3, 0.4])\n",
    "    }\n",
    "    \n",
    "    new_movie_embedding = initialize_new_movie(\n",
    "        new_movie_id, \n",
    "        similar_movies,\n",
    "        dummy_embeddings\n",
    "    )\n",
    "    print(\"\\n7. Cold start solution tested\")\n",
    "    print(f\"New movie embedding: {new_movie_embedding}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_recommendation_system()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing model training...\n",
      "\n",
      "Starting training...\n",
      "Epoch 0, Train Loss: 0.2519, Val Loss: 0.5463\n",
      "Epoch 1, Train Loss: 0.2960, Val Loss: 0.5100\n",
      "Epoch 2, Train Loss: 0.3089, Val Loss: 0.4765\n",
      "Training completed successfully!\n",
      "Model checkpoint saved successfully\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/0t/6htcr5yd28936ddkd05m8kyc0000gn/T/ipykernel_34871/652245114.py:5: UserWarning: Using a target size (torch.Size([32, 1])) that is different to the input size (torch.Size([1, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.\n",
      "  return F.mse_loss(predictions, targets)\n",
      "/var/folders/0t/6htcr5yd28936ddkd05m8kyc0000gn/T/ipykernel_34871/3274924535.py:50: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  model.load_state_dict(torch.load('best_model.pth'))\n"
     ]
    }
   ],
   "source": [
    "def test_training():\n",
    "    print(\"Testing model training...\")\n",
    "    \n",
    "    # Create a simple GNN model for testing\n",
    "    class SimpleGNN(nn.Module):\n",
    "        def __init__(self, input_dim, hidden_dim):\n",
    "            super(SimpleGNN, self).__init__()\n",
    "            self.layer1 = nn.Linear(input_dim, hidden_dim)\n",
    "            self.layer2 = nn.Linear(hidden_dim, 1)\n",
    "\n",
    "        def forward(self, graph, user_nodes, pos_movies, neg_movies):\n",
    "            # Simplified forward pass for testing\n",
    "            batch_size = len(user_nodes)\n",
    "            device = next(self.parameters()).device\n",
    "            \n",
    "            # Create tensors with proper shape and gradient tracking\n",
    "            x = torch.randn(batch_size, self.layer1.in_features, device=device, requires_grad=True)\n",
    "            \n",
    "            # Generate predictions using the same input features\n",
    "            pred_ratings = self.layer2(F.relu(self.layer1(x)))\n",
    "            pos_scores = self.layer2(F.relu(self.layer1(x)))\n",
    "            neg_scores = self.layer2(F.relu(self.layer1(x)))\n",
    "            \n",
    "            return pred_ratings, pos_scores, neg_scores\n",
    "\n",
    "    # Create sample data\n",
    "    interactions, user_data = create_sample_data()\n",
    "    \n",
    "    # Build and prepare graphs\n",
    "    train_graph = build_movie_graph(interactions[:4])  # 80% for training\n",
    "    val_graph = build_movie_graph(interactions[4:])    # 20% for validation\n",
    "    \n",
    "    # Initialize model and optimization components\n",
    "    input_dim = 10  # Example dimension\n",
    "    hidden_dim = 32\n",
    "    model = SimpleGNN(input_dim, hidden_dim)\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n",
    "        optimizer, mode='min', factor=0.5, patience=2, verbose=True\n",
    "    )\n",
    "    \n",
    "    # Define batch generation function for testing\n",
    "    def generate_batches(graph, batch_size=2):\n",
    "        for _ in range(3):  # Generate 3 dummy batches\n",
    "            subgraph, user_nodes, pos_movies, neg_movies = create_mini_batch(\n",
    "                graph, batch_size=batch_size, n_pos=2, n_neg=2,\n",
    "                n_neighbors=2, n_hops=1  # Add required parameters\n",
    "            )\n",
    "            ratings = torch.rand(batch_size, 1, requires_grad=True)  # Add requires_grad\n",
    "            yield user_nodes, pos_movies, neg_movies, ratings\n",
    "    \n",
    "    # Test training\n",
    "    print(\"\\nStarting training...\")\n",
    "    try:\n",
    "        trained_model = train_model(\n",
    "            model=model,\n",
    "            train_graph=train_graph,\n",
    "            val_graph=val_graph,\n",
    "            optimizer=optimizer,\n",
    "            scheduler=scheduler,\n",
    "            epochs=3,\n",
    "            patience=2\n",
    "        )\n",
    "        print(\"Training completed successfully!\")\n",
    "        \n",
    "        # Test model saving/loading\n",
    "        if os.path.exists('best_model.pth'):\n",
    "            print(\"Model checkpoint saved successfully\")\n",
    "        return trained_model, val_graph\n",
    "            \n",
    "    except Exception as e:\n",
    "        print(f\"Training failed with error: {str(e)}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    trained_model, val_graph = test_training()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing model predictions...\n",
      "\n",
      "Test Results:\n",
      "Number of test users: 1\n",
      "Predicted ratings shape: torch.Size([1, 1])\n",
      "Sample predicted rating: 0.105\n",
      "Sample actual rating: 0.944\n",
      "Positive score: 0.105\n",
      "Negative score: 0.105\n",
      "\n",
      "Test Loss: 0.5812\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/0t/6htcr5yd28936ddkd05m8kyc0000gn/T/ipykernel_34871/652245114.py:5: UserWarning: Using a target size (torch.Size([2, 1])) that is different to the input size (torch.Size([1, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.\n",
      "  return F.mse_loss(predictions, targets)\n"
     ]
    }
   ],
   "source": [
    "def test_model_predictions(trained_model, test_graph):\n",
    "    print(\"Testing model predictions...\")\n",
    "    \n",
    "    try:\n",
    "        # Set model to evaluation mode\n",
    "        trained_model.eval()\n",
    "        \n",
    "        # Generate test batch\n",
    "        batch = next(generate_batches(test_graph, batch_size=2))\n",
    "        user_nodes, pos_movies, neg_movies, actual_ratings = batch\n",
    "        \n",
    "        # Get predictions\n",
    "        with torch.no_grad():\n",
    "            pred_ratings, pos_scores, neg_scores = trained_model(\n",
    "                test_graph, \n",
    "                user_nodes, \n",
    "                pos_movies, \n",
    "                neg_movies\n",
    "            )\n",
    "        \n",
    "        # Print results\n",
    "        print(\"\\nTest Results:\")\n",
    "        print(f\"Number of test users: {len(user_nodes)}\")\n",
    "        print(f\"Predicted ratings shape: {pred_ratings.shape}\")\n",
    "        print(f\"Sample predicted rating: {pred_ratings[0].item():.3f}\")\n",
    "        print(f\"Sample actual rating: {actual_ratings[0].item():.3f}\")\n",
    "        print(f\"Positive score: {pos_scores[0].item():.3f}\")\n",
    "        print(f\"Negative score: {neg_scores[0].item():.3f}\")\n",
    "        \n",
    "        # Calculate test loss\n",
    "        test_loss = combined_loss(pred_ratings, actual_ratings, pos_scores, neg_scores)\n",
    "        print(f\"\\nTest Loss: {test_loss.item():.4f}\")\n",
    "        \n",
    "        return test_loss.item()\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"Testing failed with error: {str(e)}\")\n",
    "        return None\n",
    "\n",
    "# Add this to the test_training() function after successful training:\n",
    "if trained_model is not None:\n",
    "    test_model_predictions(trained_model, val_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
